{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d6348096-b55b-42d8-8bad-ba34ab9ff974",
   "metadata": {},
   "source": [
    "# QNN for the XOR Problem"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcd6b805-f7c6-4a33-b9f3-17f267c20237",
   "metadata": {},
   "source": [
    "Classiq has an available dataset for training a PQC (parameterized quantum circuit) to imitate the XOR gate, similar to how we trained a U-gate to act as a NOT gate. Design a QNN to solve the XOR problem. Read more on the dataset [here](https://docs.classiq.io/latest/reference-manual/built-in-algorithms/qml/qnn/datasets/#datasetxor)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b8a21e4-588a-4117-a6a3-323bc81b141e",
   "metadata": {},
   "source": [
    "Let's remember that the behavior of the `XOR` gate is as follows:\n",
    "\n",
    "|b0|b1|XOR|\n",
    "|:--:|:--:|:--:|\n",
    "|0|0|0|\n",
    "|0|1|1|\n",
    "|1|0|1|\n",
    "|1|1|0|\n",
    "\n",
    "Where `b0` and `b1` are bits.\n",
    "\n",
    "Now, we can get the same behavior in a quantum circuit using the `CNOT` quantum gate as follows:\n",
    "\n",
    "$$CNOT|00\\rangle = |00\\rangle$$\n",
    "$$CNOT|01\\rangle = |01\\rangle$$\n",
    "$$CNOT|10\\rangle = |11\\rangle$$\n",
    "$$CNOT|11\\rangle = |10\\rangle$$\n",
    "\n",
    "The left qubit is the control, and the right qubit is the target. We can observe that we get the result of the `XOR` operator in the target qubit (the right one)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e42c535-2195-45b4-9cc2-eaf13bf150e3",
   "metadata": {},
   "source": [
    "Knowing this, I adjusted a parameter within the `CRX` gate so that we would ultimately determine the `CNOT` gate; that is, the parameter would be equal to $\\pi$.\n",
    "\n",
    "It should be noted that the equivalent of gates can be used:\n",
    "\n",
    "$$CNOT = (I\\otimes H) \\, (CZ) \\, (I\\otimes H)$$\n",
    "\n",
    "With these gates, the desired behavior is also found, but it takes more iterations, and more oscillations were observed in the calculation of the parameter, which is why the code is left with the `CRX` gate.\n",
    "\n",
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b00e6fd0-abc6-427d-b3e4-45f3a8432ce7",
   "metadata": {},
   "source": [
    "First let's see what the Dataset for the `XOR` operation is like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d63f72a0-29b1-43bb-a5ae-04a60c269581",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--> Data for training:\n",
      "tensor([[0., 0.],\n",
      "        [1., 1.],\n",
      "        [0., 1.],\n",
      "        [1., 0.]])\n",
      "--> Corresponding labels:\n",
      "tensor([0., 0., 1., 1.])\n"
     ]
    }
   ],
   "source": [
    "from classiq.applications.qnn.datasets import DATALOADER_XOR\n",
    "\n",
    "for data, label in DATALOADER_XOR:\n",
    "    print(f\"--> Data for training:\\n{data}\")\n",
    "    print(f\"--> Corresponding labels:\\n{label}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b3ff014-69bb-4356-8b90-9c65139ac1a4",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "We can observe an important difference with respect to the dataset for the `NOT`, that in this case, the inputs are indicated as states `0` and `1`, instead of an angle that can be used in the `CRX` gate to embed the inputs in the quantum circuit.\n",
    "\n",
    "Knowing that, we can still continue using the `CRX` gate, since we will only need to encode the states `0` and `1`, so we take that information that the dataset will provide and multiply it by $\\pi$, so that we will have the corresponding state after performing the encoding (`encoding` method).\n",
    "\n",
    "And for what was previously mentioned, in the `mixing` method a `CRX` gate will be used, with a single parameter (`angle`) to be determined by the Quantum Neural Network.\n",
    "\n",
    "<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dced17d8-bb8f-47c9-b185-c14ce043df01",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opening: https://platform.classiq.io/circuit/9966091d-2f51-4a57-bda1-377af5a4c35e?version=0.42.1\n"
     ]
    }
   ],
   "source": [
    "from classiq import synthesize, qfunc, QArray, H, RX, CRX, CRZ, Output, allocate, CInt, CReal, QBit, show\n",
    "from classiq.qmod.quantum_function import create_model\n",
    "import numpy as np\n",
    "\n",
    "@qfunc\n",
    "def encoding(state0: CInt, state1: CInt, q: QArray[QBit]) -> None:\n",
    "    RX(theta=state0*np.pi, target=q[0])\n",
    "    RX(theta=state1*np.pi, target=q[1])\n",
    "\n",
    "\n",
    "@qfunc\n",
    "def mixing(theta: CReal, q: QArray[QBit]) -> None:\n",
    "    \n",
    "    #H(q[0])                   # these three gates are equivalent to CRX,\n",
    "    #CRZ(theta, q[1], q[0])    #      this option also works,\n",
    "    #H(q[0])                   #      but takes more iterations\n",
    "    \n",
    "    CRX(theta, q[1], q[0])\n",
    "\n",
    "\n",
    "@qfunc\n",
    "def main(input_0: CInt, input_1: CInt, weight_0: CReal, res: Output[QArray[QBit]]) -> None:\n",
    "    allocate(2, res)\n",
    "    encoding(state0=input_0, state1=input_1, q=res)   # loading input (two values, four possible combinations): 00, 01, 10, 11 \n",
    "    mixing(theta=weight_0, q=res)                     # adjustable parameter (an angle)\n",
    "\n",
    "\n",
    "model = create_model(main)\n",
    "\n",
    "quantum_program = synthesize(model)\n",
    "show(quantum_program)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8a959d1-e094-4d85-abc5-b0bff1f43a34",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "We define the method that will be responsible for calling `execute_qnn`, given the simplicity of this QNN we do not need more configuration.\n",
    "\n",
    "Then the method that will do the post-processing is defined, so that comparisons of what is measured from the quantum circuit can be carried out against the expected labels according to what is specified by the `XOR` dataset provided.\n",
    "\n",
    "Here we note that states `01` and `10` should give us a `1` as a result of the `XOR`, and since we are applying a gate that we want to behave like a `CNOT`, then the outputs must be `01` and `11` respectively, and for the other two states we want an output equal to `0` (the result of the `XOR`).\n",
    "\n",
    "<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c5a48c01-2236-47e2-b4b8-f5be40c42bce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from classiq.applications.qnn.types import MultipleArguments, SavedResult, ResultsCollection\n",
    "\n",
    "\n",
    "from classiq.execution import execute_qnn\n",
    "from classiq.synthesis import SerializedQuantumProgram\n",
    "\n",
    "\n",
    "def execute(\n",
    "    quantum_program: SerializedQuantumProgram, arguments: MultipleArguments\n",
    ") -> ResultsCollection:\n",
    "    return execute_qnn(quantum_program, arguments)\n",
    "\n",
    "\n",
    "# Post-process the result\n",
    "# This function validates if output is 01 or 11 then the result\n",
    "#   of the XOR operation is '1', and since we don't want the \n",
    "#   other two (00, 10) we substract their probailities\n",
    "# The returning value will correspond to the XOR result, therefore\n",
    "#   if we calculate a negative value, then we simply return a zero\n",
    "def post_process(result: SavedResult) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Take in a `SavedResult` with `ExecutionDetails` value type, and return the\n",
    "    probability of measuring |01> + |11> - |00> - |10>\n",
    "    \"\"\"\n",
    "    counts: dict = result.value.counts\n",
    "    xor_result: float = counts.get(\"01\", 0.0) / sum(counts.values()) + counts.get(\"11\", 0.0) / sum(counts.values()) \\\n",
    "                        - counts.get(\"00\", 0.0) / sum(counts.values()) - counts.get(\"10\", 0.0) / sum(counts.values())\n",
    "    \n",
    "    return torch.tensor(0.0 if xor_result < 0 else xor_result)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f69b5d13-df32-40e4-83bc-c8f8dd71f0f3",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "Now we create a neural network, to which we add a `QLayer`, which is responsible for executing the quantum circuit.\n",
    "\n",
    "This process is the same as creating a NN with PyTorch, only we use a special layer provided by the Classiq SDK.\n",
    "\n",
    "<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "55ee6c51-771b-4f31-ae66-c4e1ab8f52e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from classiq.applications.qnn import QLayer\n",
    "\n",
    "\n",
    "class Net(torch.nn.Module):\n",
    "    def __init__(self, *args, **kwargs) -> None:\n",
    "        super().__init__()\n",
    "        self.qlayer = QLayer(\n",
    "            quantum_program,  # the quantum program, the result of `synthesize()`\n",
    "            execute,  # a callable that takes\n",
    "            # - a quantum program\n",
    "            # - parameters to that program (a tuple of dictionaries)\n",
    "            # and returns a `ResultsCollection`\n",
    "            post_process,  # a callable that takes\n",
    "            # - a single `SavedResult`\n",
    "            # and returns a `torch.Tensor`\n",
    "            *args,\n",
    "            **kwargs\n",
    "        )\n",
    "        #self.qlayer.weight.data.\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        # return the new parameter\n",
    "        return self.qlayer(x)\n",
    "\n",
    "\n",
    "model = Net()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c50d60b-d767-4337-ad13-c51918dc5a62",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "We indicate the dataset for the `XOR` operation. We continue using the _Mean Absolute Error_ to calculate how different the output of the labels is, this function works because we are calculating it that way within `post_process`. And finally we continue using _Stochastic Gradient Descent_ as an optimizer, since it is good enough to determine the desired parameter.\n",
    "\n",
    "<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bb37c29f-81a2-4408-b883-73e43378d495",
   "metadata": {},
   "outputs": [],
   "source": [
    "from classiq.applications.qnn.datasets import DATALOADER_XOR\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "_LEARNING_RATE = 1\n",
    "\n",
    "# choosing our data\n",
    "data_loader = DATALOADER_XOR     # Dataset to train the XOR operation\n",
    "# choosing our loss function\n",
    "loss_func = nn.L1Loss()          # Mean Absolute Error (MAE)\n",
    "# choosing our optimizer\n",
    "optimizer = optim.SGD(model.parameters(), lr=_LEARNING_RATE)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "722b6047-5f24-4dad-982c-74aecf906a4a",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "We train the QNN. We print the labels right next to the model output to see the progress, as well as the current parameter in the corresponding iteration (the epoch).\n",
    "\n",
    "<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2103a0ae-1ff8-4402-aeb4-c24e69e401f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 Parameter containing:\n",
      "tensor([0.7793], requires_grad=True)\n",
      "label: tensor([1., 1., 0., 0.])\n",
      "output: tensor([1.0000, 0.0000, 0.7061, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "1 Parameter containing:\n",
      "tensor([0.9746], requires_grad=True)\n",
      "label: tensor([1., 0., 0., 1.])\n",
      "output: tensor([1.0000, 0.0000, 0.5762, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "2 Parameter containing:\n",
      "tensor([1.0357], requires_grad=True)\n",
      "label: tensor([0., 1., 1., 0.])\n",
      "output: tensor([0.5244, 1.0000, 0.0000, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "3 Parameter containing:\n",
      "tensor([0.5840], requires_grad=True)\n",
      "label: tensor([1., 0., 1., 0.])\n",
      "output: tensor([0.0000, 0.8242, 1.0000, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "4 Parameter containing:\n",
      "tensor([0.6450], requires_grad=True)\n",
      "label: tensor([0., 1., 1., 0.])\n",
      "output: tensor([0.0000, 0.0000, 1.0000, 0.7949], grad_fn=<QLayerFunctionBackward>)\n",
      "5 Parameter containing:\n",
      "tensor([1.0235], requires_grad=True)\n",
      "label: tensor([0., 0., 1., 1.])\n",
      "output: tensor([0.0000, 0.4824, 1.0000, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "6 Parameter containing:\n",
      "tensor([1.4507], requires_grad=True)\n",
      "label: tensor([1., 0., 1., 0.])\n",
      "output: tensor([0.0000, 0.0869, 1.0000, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "7 Parameter containing:\n",
      "tensor([1.6094], requires_grad=True)\n",
      "label: tensor([1., 0., 0., 1.])\n",
      "output: tensor([0.0684, 0.0000, 0.0000, 1.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "8 Parameter containing:\n",
      "tensor([2.5249], requires_grad=True)\n",
      "label: tensor([0., 1., 0., 1.])\n",
      "output: tensor([0.0000, 0.8057, 0.0000, 1.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "9 Parameter containing:\n",
      "tensor([2.7935], requires_grad=True)\n",
      "label: tensor([1., 0., 1., 0.])\n",
      "output: tensor([1.0000, 0.0000, 0.9453, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "10 Parameter containing:\n",
      "tensor([2.8545], requires_grad=True)\n",
      "label: tensor([0., 1., 1., 0.])\n",
      "output: tensor([0.0000, 0.9473, 1.0000, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "11 Parameter containing:\n",
      "tensor([3.0742], requires_grad=True)\n",
      "label: tensor([0., 1., 0., 1.])\n",
      "output: tensor([0.0000, 1.0000, 0.0000, 0.9990], grad_fn=<QLayerFunctionBackward>)\n",
      "12 Parameter containing:\n",
      "tensor([3.1231], requires_grad=True)\n",
      "label: tensor([1., 1., 0., 0.])\n",
      "output: tensor([1.0000, 0.9990, 0.0000, 0.0000], grad_fn=<QLayerFunctionBackward>)\n",
      "13 Parameter containing:\n",
      "tensor([3.1353], requires_grad=True)\n",
      "label: tensor([1., 0., 1., 0.])\n",
      "output: tensor([1., 0., 1., 0.], grad_fn=<QLayerFunctionBackward>)\n",
      "14 Parameter containing:\n",
      "tensor([3.1353], requires_grad=True)\n",
      "label: tensor([1., 0., 0., 1.])\n",
      "output: tensor([1., 0., 0., 1.], grad_fn=<QLayerFunctionBackward>)\n",
      "15 Parameter containing:\n",
      "tensor([3.1353], requires_grad=True)\n",
      "label: tensor([0., 0., 1., 1.])\n",
      "output: tensor([0., 0., 1., 1.], grad_fn=<QLayerFunctionBackward>)\n",
      "16 Parameter containing:\n",
      "tensor([3.1353], requires_grad=True)\n",
      "label: tensor([0., 0., 1., 1.])\n",
      "output: tensor([0., 0., 1., 1.], grad_fn=<QLayerFunctionBackward>)\n",
      "17 Parameter containing:\n",
      "tensor([3.1353], requires_grad=True)\n",
      "label: tensor([0., 0., 1., 1.])\n",
      "output: tensor([0., 0., 1., 1.], grad_fn=<QLayerFunctionBackward>)\n",
      "18 Parameter containing:\n",
      "tensor([3.1353], requires_grad=True)\n",
      "label: tensor([1., 0., 1., 0.])\n",
      "output: tensor([1., 0., 1., 0.], grad_fn=<QLayerFunctionBackward>)\n",
      "19 Parameter containing:\n",
      "tensor([3.1353], requires_grad=True)\n",
      "label: tensor([0., 1., 0., 1.])\n",
      "output: tensor([0., 1., 0., 1.], grad_fn=<QLayerFunctionBackward>)\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "\n",
    "def train(\n",
    "    model: nn.Module,\n",
    "    data_loader: DataLoader,\n",
    "    loss_func: nn.modules.loss._Loss,\n",
    "    optimizer: optim.Optimizer,\n",
    "    epoch: int = 20,\n",
    ") -> None:\n",
    "    for index in range(epoch):\n",
    "        print(index, model.qlayer.weight)\n",
    "        for data, label in data_loader:\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            output = model(data)\n",
    "            \n",
    "            print('label:', label)     # print the expected values along side with the calculated ones\n",
    "            print('output:', output)\n",
    "\n",
    "            loss = loss_func(output, label)\n",
    "            loss.backward()\n",
    "\n",
    "            optimizer.step()\n",
    "\n",
    "\n",
    "train(model, data_loader, loss_func, optimizer)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b445fe7e-ae01-41e5-9d0a-46f2fa8e7819",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "Finally we check the accuracy, in order to know how good the model is, how well the QNN has found the desired parameter. \n",
    "\n",
    "<br>\n",
    "\n",
    "We observe that the **accuary is 100%**, so we can have the complete example for the `XOR` operation.\n",
    "\n",
    "<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f62cea58-dac5-4c72-a3af-b62536df8629",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predictions: tensor([1., 0., 0., 1.], requires_grad=True)\n",
      "labels:      tensor([1., 0., 0., 1.])\n",
      "Test Accuracy of the model: 100.00%\n"
     ]
    }
   ],
   "source": [
    "def check_accuracy(model: nn.Module, data_loader: DataLoader, atol=1e-4) -> float:\n",
    "    num_correct = 0\n",
    "    total = 0\n",
    "    model.eval()\n",
    "\n",
    "    with torch.no_grad():                        # temporarily disable gradient calculation\n",
    "        for data, labels in data_loader:\n",
    "            # let the model predict\n",
    "            predictions = model(data)\n",
    "            print('predictions:', predictions)\n",
    "            print('labels:     ', labels)\n",
    "\n",
    "            # get a tensor of booleans, indicating if each label is close to the real label\n",
    "            is_prediction_correct = predictions.isclose(labels, atol=atol)\n",
    "\n",
    "            # count the amount of `True` predictions\n",
    "            num_correct += is_prediction_correct.sum().item()\n",
    "            # count the total evaluations, the first dimension of `labels` is `batch_size`\n",
    "            total += labels.size(0)\n",
    "\n",
    "    accuracy = float(num_correct) / float(total)\n",
    "    return accuracy\n",
    "\n",
    "\n",
    "accuracy = check_accuracy(model, data_loader)\n",
    "\n",
    "print(f\"Test Accuracy of the model: {accuracy*100:.2f}%\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
